//
// Created by ECNU on 2024/2/22.
//

#include <nanofaiss/utils/Heap.h>

namespace faiss {

template<typename C>
void HeapArray<C>::heapify() {
#pragma omp parallel for
    for (size_t i = 0; i != nh; i++)
        heap_heapify<C>(k, val + i * k, ids + i * k);
}

template<typename C>
void HeapArray<C>::reorder() {
#pragma omp parallel for
    for (size_t i = 0; i != nh; i++)
        heap_reorder<C>(k, val + i * k, ids + i * k);
}

template<typename C>
void HeapArray<C>::addn(size_t nj, const T* vin, TI j0, size_t i0, int64_t ni) {
    if (ni == -1)
        ni = nh;
    FAISS_THROW_IF_NOT(i0 >=0 && i0 + ni <= nh);
#pragma omp parallel for if (nj * ni > 100000)
    for (size_t i = i0; i != i0 + ni; i++) {
        T* simi = get_val(i);
        TI* idx = get_id(i);
        const T* ip_line = vin + (i - i0) * nj;

        for (size_t j = 0; j < nj; j++) {
            T ip = ip_line[j];
            if (C::cmp(simi[0], ip)) {
                heap_replace_top<C>(k, simi, idx, ip, j + j0);
            }
        }
    }
}


template<typename C>
void HeapArray<C>::add_with_ids(size_t nj,
                                const T* vin,
                                const TI* id_in,
                                int64_t id_stride,
                                size_t i0,
                                int64_t ni) {
    if (id_in == nullptr) {
        addn(nj, vin, 0, i0, ni);
        return;
    }
    if (ni == -1)
        ni = nh;
    FAISS_THROW_IF_NOT(i0 >= 0 && i0 + ni <= nh);
#pragma omp parallel for if (ni * nj > 100000)
    for (size_t i = i0; i < i0 + ni; i++) {
        T* simi = get_val(i);
        TI* idx = get_id(i);
        const T* ip_line = vin + (i - i0) * nj;
        const TI *id_line = id_in + (i - i0) * id_stride; // actually id_stride should equal to nj

        for (size_t j = 0; j < nj; j++) {
            T ip = ip_line[j];
            if (C::cmp(simi[0], ip)) {
                heap_replace_top<C>(k, simi, idx, ip, id_line[j]);
            }
        }
    }

}

}



